package edu.arizona.sista.learn.evaluation;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;

import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public class CMUXML {
	
	private File _xmlFile;
	private int _winSize;
	private int _shiftSize;
	private String _movieName;

    private Map<String, VideoObject> _videoObjects;
    private List<VideoInterval> _videoIntervals;
    private List<List<VideoAction>> _timeline;
    
    private int _maxSliceIndex;
    private int _maxFrameIndex;
    
    class VideoObject {
    	public String UID;
    	public String type;
    	public int startFrameIndex;
    	public int endFrameIndex;
    	public int totalPoints;
    	VideoObject(String uid, String type) {
    		this.UID = uid; this.type = type;
    	}
    	VideoObject(String uid, String type, int startIndex, int endIndex, int totalPoints) {
    		this.UID = uid; this.type = type;
    		this.startFrameIndex = startIndex;
    		this.endFrameIndex = endIndex;
    		this.totalPoints = totalPoints;
    	}
    }
    
    class VideoAction {
    	public String verb;
    	public String actor, actee, acton;
    	VideoAction(String verb, String actor, String actee, String acton) {
    		this.verb = verb; this.actor = actor; this.actee = actee; this.acton = acton;
    	}
    }
    
    class VideoInterval {
    	public VideoAction action;
    	public int startFrameIndex, endFrameIndex;
    	public int startSliceIndex, endSliceIndex;
    	VideoInterval(String verb, String actor, String actee, String acton, int start, int end) {
    		this.action = new VideoAction(verb, actor, actee, acton);
    		this.startFrameIndex = start; this.endFrameIndex = end;
    		this.startSliceIndex = start / _shiftSize;
        	this.endSliceIndex = end / _shiftSize;
    	}
    	public boolean frameOverlaps(int startIndex, int endIndex) {
    		return ((startIndex <= this.endFrameIndex) && (endIndex >= this.startFrameIndex));
    	}
    	public boolean sliceOverlaps(int startIndex, int endIndex) {
    		return ((startIndex <= this.endSliceIndex) && (endIndex >= this.startSliceIndex));
    	}
    }
	
	public CMUXML(File xml) {
		this(xml, 1, 1);
	}
	
	public CMUXML(File xml, int win, int shift) { 
		_xmlFile = xml;
		_winSize = win;
		_shiftSize = shift;
		_movieName = xml.getName();
	}
	
	private String getTextAttribute(Element ele, String key) {
		return ele.getAttribute(key);
	}

	private int getIntAttribute(Element ele, String key) {
		try {
			return Integer.parseInt(getTextAttribute(ele,key));
		} catch (Exception e) {
			e.printStackTrace();
		}
		return Integer.MIN_VALUE;
	}
	
	public void parse() {
		parse(1);
	}
	
	public void parse(int mode) {
		try {
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			DocumentBuilder db = dbf.newDocumentBuilder();
			Document dom = db.parse(_xmlFile);
			
			// normalize text representation
			Element root = dom.getDocumentElement();
			Element annotationData = (Element)root.getElementsByTagName("VideoAnnotationData").item(0);
			Element annotationActions = (Element)root.getElementsByTagName("VideoAnnotationActions").item(0);
			
            // Get max frame
			if (_movieName.startsWith("HAND1_A1_C2_Act6_2_URBAN_MC_AFTN_b460e682-07b6-11e0-9fc5-e80688cb869a")) {
				int a = 2;
			}
            _maxFrameIndex = Math.max(0, getIntAttribute(annotationData, "frames") - 1);
            _maxSliceIndex = (_maxFrameIndex - _winSize) / _shiftSize;
            
            // Grab object types & ID's
            _videoObjects = new HashMap<String, VideoObject>();
			Element trackObjects = (Element) annotationData.getElementsByTagName("trackObjects").item(0);
			NodeList objects = trackObjects.getElementsByTagName("trackObject");
            for (int i = 0; i < objects.getLength(); i++) {
            	Element obj = (Element) objects.item(i);
            	String uid = getTextAttribute(obj, "ObjectUID");
            	String type = getTextAttribute(obj, "type");
            	int totalPoints = getIntAttribute(obj, "points");
            	NodeList points = obj.getElementsByTagName("trackPoint");
            	assert(totalPoints == points.getLength());
            	if (points.getLength() > 0) {
	            	Element firstPoint = (Element) points.item(0);
	            	Element endPoint = (Element) points.item(points.getLength()-1);
	            	int startIndex = Math.max(0, getIntAttribute(firstPoint, "Frame") - 1);
	            	int endIndex = Math.max(0, getIntAttribute(endPoint, "Frame") - 1);
	            	_videoObjects.put(uid, new VideoObject(uid, type, startIndex, endIndex, totalPoints));
	            	
	            } else {
	            	_videoObjects.put(uid, new VideoObject(uid, type, -1, -1, totalPoints));
	            }
            	
            }
            
            // Grab intervals
            _videoIntervals = new ArrayList<VideoInterval>();
            Element verbRelations = (Element)annotationActions.getElementsByTagName("verbRelations").item(0);
            NodeList actions = verbRelations.getElementsByTagName("verbRelation");
            for (int i = 0; i < actions.getLength(); i++) {
            	Element action = (Element)actions.item(i);
            	String verb = getTextAttribute(action, "verb");
            	String actor = getTextAttribute(action, "actor");
            	String actee = getTextAttribute(action, "actee");
            	String acton = getTextAttribute(action, "acton");
            	int startIndex = Math.max(0, getIntAttribute(action, "frameS") - 1);
            	int endIndex = Math.max(0, getIntAttribute(action, "frameE") - 1);
            	
            	VideoInterval intv = new VideoInterval(verb, actor, actee, acton, startIndex, endIndex);
            	_videoIntervals.add(intv);
            }
            
            // Store intervals in a timeline format for easy look-up
            if (mode == 1)
            	constructTimeline1();
            else
            	constructTimeline2();
            
		}catch(ParserConfigurationException pce) {
			pce.printStackTrace();
		}catch(SAXException se) {
			se.printStackTrace();
		}catch(IOException ioe) {
			ioe.printStackTrace();
		}
	}
	
	public boolean verifyData() {
		boolean pass = true;
		
		System.out.println("Checking: " + _movieName);
		
		// Check objects length
		for (String k : _videoObjects.keySet()) {
			VideoObject vo = _videoObjects.get(k);
			if (vo.totalPoints == 0) {
				System.out.println("OBJ " + vo.UID + ": Empty track (0 point).");
				pass = false;
			} else if ((vo.endFrameIndex - vo.startFrameIndex + 1) != vo.totalPoints) {
	    		System.out.println("OBJ " + vo.UID + ": Fragmented track. " +
	    				"total points = " + vo.totalPoints + " | " +
	    				"first frame = " + (vo.startFrameIndex+1) + " | " +
	    				"last frame = " + (vo.endFrameIndex+1));
	    		// Not really errors, just informative
	    	}
		}
		
		// Check video intervals
		for (VideoInterval vi : _videoIntervals) {
			if (vi.startFrameIndex == vi.endFrameIndex) {
				System.out.println("ACTION: EQUAL START & END. " +
						"verb = " + vi.action.verb + " | " +
						"start = " + (vi.startFrameIndex+1) + " | " +
						"end = " + (vi.endFrameIndex+1) + " | " +
						"actor = " + getVideoObjectType(vi.action.actor) + " | " +
						"actee = " + getVideoObjectType(vi.action.actee) + " | " +
						"acton = " + getVideoObjectType(vi.action.acton));
	    		pass = false;
			}
			if (vi.startFrameIndex > vi.endFrameIndex) {
				System.out.println("ACTION: START AFTER END. " +
						"verb = " + vi.action.verb + " | " +
						"start = " + (vi.startFrameIndex+1) + " | " +
						"end = " + (vi.endFrameIndex+1) + " | " +
						"actor = " + getVideoObjectType(vi.action.actor) + " | " +
						"actee = " + getVideoObjectType(vi.action.actee) + " | " +
						"acton = " + getVideoObjectType(vi.action.acton));
	    		pass = false;
			}
			if ( (!_videoObjects.containsKey(vi.action.actor) && !"-1".equals(vi.action.actor)) ||
					(!_videoObjects.containsKey(vi.action.actee) && !"-1".equals(vi.action.actee)) ||
					(!_videoObjects.containsKey(vi.action.acton) && !"-1".equals(vi.action.acton)) ) {
				System.out.println("ACTION: UNRECOGNIZABLE actor/actee/acton. " +
						"verb = " + vi.action.verb + " | " +
						"start = " + (vi.startFrameIndex+1) + " | " +
						"end = " + (vi.endFrameIndex+1) + " | " +
						"actor = " + getVideoObjectType(vi.action.actor) + " | " +
						"actee = " + getVideoObjectType(vi.action.actee) + " | " +
						"acton = " + getVideoObjectType(vi.action.acton));

	    		pass = false;
			} else {
				if (!"-1".equals(vi.action.actor)) {
					VideoObject actor = _videoObjects.get(vi.action.actor);
					if (vi.startFrameIndex < actor.startFrameIndex) {
						System.out.println("ACTION: VERB STARTED BEFORE actor STARTED. " +
								"verb = " + vi.action.verb + " | " +
								"actor = " + getVideoObjectType(vi.action.actor) + " | " +
								"verbS = " + (vi.startFrameIndex+1) + " | " +
								"actorS = " + (actor.startFrameIndex+1));
			    		pass = false;
					}
					if (vi.endFrameIndex > actor.endFrameIndex) {
						System.out.println("ACTION: VERB ENDED AFTER actor ENDED. " +
								"verb = " + vi.action.verb + " | " +
								"actor = " + getVideoObjectType(vi.action.actor) + " | " +
								"verbE = " + (vi.endFrameIndex+1) + " | " +
								"actorE = " + (actor.endFrameIndex+1) + " | ");
			    		pass = false;
					}
				}
				if (!"-1".equals(vi.action.actee)) {
					VideoObject actee = _videoObjects.get(vi.action.actee);
					if (vi.startFrameIndex < actee.startFrameIndex) {
						System.out.println("ACTION: VERB STARTED BEFORE actee STARTED. " +
								"verb = " + vi.action.verb + " | " +
								"actee = " + getVideoObjectType(vi.action.actee) + " | " +
								"verbS = " + (vi.startFrameIndex+1) + " | " +
								"acteeS = " + (actee.startFrameIndex+1));
			    		pass = false;
					}
					if (vi.endFrameIndex > actee.endFrameIndex) {
						System.out.println("ACTION: VERB ENDED AFTER actee ENDED. " +
								"verb = " + vi.action.verb + " | " +
								"actee = " + getVideoObjectType(vi.action.actee) + " | " +
								"verbE = " + (vi.endFrameIndex+1) + " | " +
								"acteeE = " + (actee.endFrameIndex+1) + " | ");
			    		pass = false;
					}
				}
				if (!"-1".equals(vi.action.acton)) {
					VideoObject acton = _videoObjects.get(vi.action.acton);
					if (vi.startFrameIndex < acton.startFrameIndex) {
						System.out.println("ACTION: VERB STARTED BEFORE acton STARTED. " +
								"verb = " + vi.action.verb + " | " +
								"acton = " + getVideoObjectType(vi.action.acton) + " | " +
								"verbS = " + (vi.startFrameIndex+1) + " | " +
								"actonS = " + (acton.startFrameIndex+1));
			    		pass = false;
					}
					if (vi.endFrameIndex > acton.endFrameIndex) {
						System.out.println("ACTION: VERB ENDED AFTER acton ENDED. " +
								"verb = " + vi.action.verb + " | " +
								"acton = " + getVideoObjectType(vi.action.acton) + " | " +
								"verbE = " + (vi.endFrameIndex+1) + " | " +
								"actonE = " + (acton.endFrameIndex+1) + " | ");
			    		pass = false;
					}
				}
			}
		}
		
		System.out.println("");
		
		return pass;
	}
	
	private String getVideoObjectType(String id) {
		if (_videoObjects.containsKey(id))
			return _videoObjects.get(id).type;
		return id;
	}
	
	/**
	 * Each slice = shift-size number of frames.
	 */
	private void constructTimeline1() {
    	_timeline = new ArrayList<List<VideoAction>>();
    	for (int i = 0; i <= _maxSliceIndex; i++) {
    		_timeline.add(i, null);		// initialize each time tick to null
    	}
    	for (VideoInterval vi : _videoIntervals) {
    		assert(vi.endSliceIndex <= _maxSliceIndex + _winSize/_shiftSize) :
    			"Incorrect ending slice index: " + vi.action.verb + " " + vi.endFrameIndex +
    			" " + vi.endSliceIndex + " " + _maxSliceIndex;
    		for (int t = vi.startSliceIndex; t <= Math.min(vi.endSliceIndex, _maxSliceIndex); t++) {
    			if (_timeline.get(t) == null)
    				_timeline.set(t, new ArrayList<VideoAction>());
    			_timeline.get(t).add(vi.action);
    		}
    	}
	}
	
	/**
	 * Each slice = win-size number of frames (shift by shift-size)
	 */
	private void constructTimeline2() {
    	_timeline = new ArrayList<List<VideoAction>>();
    	for (int i = 0; i <= _maxSliceIndex; i++) {
    		_timeline.add(i, null);		// initialize each time tick to null
    		
    		int startFrameIndx = i * _shiftSize;
    		int endFrameIndx = startFrameIndx + _winSize - 1;
    		for (VideoInterval vi : _videoIntervals) {
    			if (vi.frameOverlaps(startFrameIndx, endFrameIndx)) {
        			if (_timeline.get(i) == null)
        				_timeline.set(i, new ArrayList<VideoAction>());
        			_timeline.get(i).add(vi.action);
    			}
    		}
    	}
	}

	/**
	 * Get list of actions/verbs that occurred at time t.
	 * @param t Time tick.
	 * @return List of actions/verbs that occurred at time t.
	 */
	public List<VideoAction> getActions(int t) {
		if (t >= _timeline.size())
			return null;
		return _timeline.get(t);
	}
	
	/**
	 * Return true or false whether the given verb occurred at the
	 * specified time.
	 * @param verb Verb to check.
	 * @param t Time tick.
	 * @return True if the given verb occurred at the specified time.
	 */
	public boolean containsAt(String verb, int t) {
		if (t >= _timeline.size())
			return false;
		List<VideoAction> actions = _timeline.get(t);
		if (actions != null) {
			for (VideoAction a : actions) {
				if (a.verb.equalsIgnoreCase(verb)) {
					return true;
				}
			}
		}
		return false;
	}
	
	/**
	 * Return true or false whether the given verb with some specific
	 * type binding occurred at the specified time.
	 * @param verb Verb to check.
	 * @param t Time tick.
	 * @param type List of type restrictions in lower case.
	 * @return True if the given bounded-verb occurred at the specified time.
	 */
	public boolean containsAt(String verb, int t, Set<String> type) {
		if (t >= _timeline.size())
			return false;
		List<VideoAction> actions = _timeline.get(t);
		if (actions != null) {
			for (VideoAction a : actions) {
				if (a.verb.equalsIgnoreCase(verb)) {
					if (!a.actee.equals("-1") && !type.contains(getVideoObjectType(a.actee).toLowerCase()))
						continue;
					if (!a.acton.equals("-1") && !type.contains(getVideoObjectType(a.acton).toLowerCase()))
						continue;
					if (!a.actor.equals("-1") && !type.contains(getVideoObjectType(a.actor).toLowerCase()))
						continue;
					return true;
				}
			}
		}
		return false;
	}
	
	public List<List<VideoAction>> getTimeline() {
		return _timeline;
	}
	
	public int getMovieLength() {
		assert(_timeline.size() == _maxSliceIndex+1);
		return _maxSliceIndex + 1;
	}
}
